package tt.langchain4j.rag;

import dev.langchain4j.data.message.UserMessage;
import dev.langchain4j.rag.AugmentationRequest;
import dev.langchain4j.rag.AugmentationResult;
import dev.langchain4j.rag.RetrievalAugmentor;
import dev.langchain4j.rag.query.Metadata;

import java.util.Arrays;
import java.util.List;

public class MultiableRetribevalAugmentor implements RetrievalAugmentor {


    private final List<RetrievalAugmentor> retrievalAugmentorList;

    public MultiableRetribevalAugmentor(List<RetrievalAugmentor> retrievalAugmentorList) {
        this.retrievalAugmentorList = retrievalAugmentorList;
    }


    public static MultiableRetribevalAugmentor from(RetrievalAugmentor... retrievalAugmentor) {
        if ( null == retrievalAugmentor || retrievalAugmentor.length == 0 ) {
            throw new RuntimeException("must no-empty");
        }
        return new MultiableRetribevalAugmentor(Arrays.asList(retrievalAugmentor));
    }


    @Override
    public UserMessage augment(UserMessage userMessage, Metadata metadata) {
        AugmentationRequest augmentationRequest = new AugmentationRequest(userMessage, metadata);
        return (UserMessage) augment(augmentationRequest).chatMessage();
    }

    @Override
    public AugmentationResult augment(AugmentationRequest request) {
        AugmentationResult result = null;
        for (RetrievalAugmentor retrievalAugmentor : retrievalAugmentorList) {
            result = retrievalAugmentor.augment(request);
            //继续！！！！
            request = new AugmentationRequest(result.chatMessage(),  request.metadata() );
        }
        return result;
    }
}
